import pdfplumber
from document.document import Document
from document.page import Page
from document.content import Content, ContentType, TableContent
from utils.logger import LOG

class PDFParser:
    ts = None

    # Table settings.
    def __init__(self):
        self.ts = {
            "vertical_strategy": "lines",
            "horizontal_strategy": "lines",
        }

    def parse(self, pdf_file):
        document = Document()
        with pdfplumber.open(pdf_file) as pdf:
            for p in pdf.pages:
                page = Page()
                # 提取PDF的文本信息，并过滤掉表格数据
                boxes = [table.bbox for table in p.find_tables(table_settings=self.ts)]
                raw_text = p.filter(lambda obj: self.not_within_bboxes(obj, boxes)).extract_text()
                tables = p.find_tables()

                # Handling text
                if raw_text:
                    # Remove empty lines and leading/trailing whitespaces
                    raw_text_lines = raw_text.splitlines()
                    cleaned_raw_text_lines = [line.strip() for line in raw_text_lines if line.strip()]
                    cleaned_raw_text = "\n".join(cleaned_raw_text_lines)

                    text_content = Content(content_type=ContentType.TEXT, original=cleaned_raw_text)
                    page.add_content(text_content)
                    LOG.debug(f"[raw_text]\n {cleaned_raw_text}")

                # Handling tables
                if tables:
                    table = TableContent(tables)
                    page.add_content(table)
                    LOG.debug(f"[table]\n{table}")

                document.add_page(page)
        return document


    # 去除Page中的表格
    def not_within_bboxes(self, obj, boxes):
        """Check if the object is in any of the table's bbox."""
        def obj_in_bbox(_bbox):
            v_mid = (obj["top"] + obj["bottom"]) / 2
            h_mid = (obj["x0"] + obj["x1"]) / 2
            x0, top, x1, bottom = _bbox
            return (h_mid >= x0) and (h_mid < x1) and (v_mid >= top) and (v_mid < bottom)

        return not any(obj_in_bbox(__bbox) for __bbox in boxes)


if __name__ == "__main__":
    PDFParser().parse('d:/dd.pdf')
